BPM: Byte-Prefix Marginalization for Cross-Tokenizer OPD
An engineering solution for cross-tokenizer on-policy distillation. The official implementation is very complicated in detail, as it needs to handle many special and edge cases.
Table of Contents
- 1. Three Targets
- 2. Solution
- 2.1. Notations
- 2.2. Cases
- 2.2.1. Base Case: Student and Teacher Token Starts at Same Position
- 2.2.2. Special Case 1: Student Token Starts in the Middle of a Teacher Token
- 2.2.3. Special Case 2: Student Tokens Spans Multiple Teacher Tokens
- 2.2.4. Special Case 3: Whitespace & Tabs Only
- 2.2.5. Special Case 4: Stop Tokens and Bridging
- 2.3. Training Objective
- 3. Findings
1. Three Targets
A well-designed tokenizer should have these 3 properties.
- Vocabulary Completeness
- Each token in student’s vocabulary must be assigned a probability score by the teacher.
- Byte-level Alignment
- The probability of a teacher’s token must be assigned to a student’s token that matches the teacher’s.
- Mass Preservation
- The sum of probability scores of all tokens in student’s vocabulary must be \(100\%\).
2. Solution
Consider that in OPD, the student emits a trajectory. The key observation is that, no matter what the tokenizer is, the underlying bytes sequence is always the same. Suppose we tokenize the string using teacher’s tokenizer and student’s tokenizer, we define
- Sync Point
- Sync points refer to a position at which both teacher’s and student’s tokenizers split. This indicates that, the substrings and the byte sequences before this sync point are the same, regardless of whether the original string is tokenized by teacher’s or student’s tokenizer. This also indicates that, teacher and student have the same context at the sync point.
- Chunk
- The substring between 2 sync points.
We just need to handle chunks properly to provide dense rewards.
2.1. Notations
2.1.1. Token Mapping
First, we need to prebuild a mapping \(\phi:\mathcal{V}^{T}\mapsto \mathcal{V}^{S}\) 1 that maps a teacher token \(v^{T} \in \mathcal{V}^{T}\) to some student token \(v^{S} \in \mathcal{V}^{S}\) such that
- \(v^{S}\) has the longest byte string and;
- the byte sequence of \(v^{S}\) is a prefix of the byte sequence of \(v^{T}\).
If we cannot find such student token, we write \(\phi(v^{T})=\varnothing\). In general,
\[ \phi(v^{T}) = \argmax_{\begin{subarray}{c}v^{S} \in \mathcal{V}^{S}\\\texttt{bytes}(v^{S}) \text{ prefixes } \texttt{bytes}(v^{T}) \end{subarray}} |\texttt{bytes}(v^{S})| \]
2.1.2. Probabililty Notation
At position \(i\), given all history token \(v_{\lt i}^{T}\), the probability of a token \(v^{T}\) that is assigned by the teacher is given by
\[ q_{i}^{T}(v^{T}) \]
2.2. Cases
The core solution of BPM is based on branching on different cases, including general cases, and various edge cases. We’ll go through them, explain what each case means, and corresponding solution.
2.2.1. Base Case: Student and Teacher Token Starts at Same Position
Suppose we are now at a sync point at position \(i\) for student and position \(T(i)\) for teacher. At this position, teacher assigns a probability \(q_{T(i)}(v^{T})\) to every token \(v^{T}\) that belongs to its vocabulary.
To assign a probability (produced by the teacher) over student’s vocabulary, for a student token \(v^{S}\), we accumulate the probability of all teacher tokens \(v^{T}\) who are prefixed by this student token (in terms of byte string). We denote this probability by \(p_{i}(v^{S})\). In case that some teacher tokens cannot be mapped to any student tokens, we accumulate their probabilities to a residual term \(p_{i}(\varnothing)\).
\[ p_{i}(v^{S}) = \sum_{\phi(v^{T})=v^{S}} q_{T(i)}(v^{T}), \qquad p_{i}(\varnothing) = \sum_{\phi(v^{T})=\varnothing} q_{T(i)}(v^{T}) \]
This residual term ensures that the mass sums to one
2.2.2. Special Case 1: Student Token Starts in the Middle of a Teacher Token
When a student token starts in the middle of a teacher token, it implies that the student has already generated a token that is the prefix of the teacher token (in terms of byte string). Let’s denote this student token by \(v^{S}_{pre}\).
The idea here is to use conditional probability and normalization to remove the effect of this generated prefix. We define \(\pi^{T}(c)\) as the sum of probabilities of teacher tokens who start with \(c\) in terms of byte string.
\[ \pi^{T}(c) = \sum_{\texttt{bytes}(c)\text{ prefixes }\texttt{bytes}(v^{T})} q_{T(i)}(v^{T}) \]
We further explain \(\phi(v^{T} | v^{S}_{pre})=v^{S}\)2: the subscript means teacher token \(v^{T}\) should first begin with \(v^{S}_{pre}\), and then followed by \(v^{S}\), therefore, equivalently, it’s equal to teacher token \(v^{T}\) is prefixed by \(v^{S}_{pre}\) concatenated by \(v^{S}\).
Then, we can define the distribution in this case, the condition notation in \(p_{i}(\cdot | v^{S}_{pre})\) suggests the student has already generated a student token \(v^{S}_{pre}\)
\[ p_{i}(v^{S} | v^{S}_{pre}) = \frac{\sum\limits_{\phi(v^{T} | v^{S}_{pre}) = v^{S}} q_{T(i)}(v^{T})}{\pi(v^{S}_{pre})} \]
2.2.3. Special Case 2: Student Tokens Spans Multiple Teacher Tokens
Suppose a student token \(v^{S}\) spans multiple teacher tokens \(v^{T}_{j}, v^{T}_{j+1}, \dots, v^{T}_{j+m-1}\) plus a trailing suffix string starting from teacher’s position \(j+m\): \(v^{S}_{suf}\). In this case, we use chain rule to handle this. In particular, we directly multiply their probabilities and \(\pi_{j+m}(v^{S}_{suf})\).
\[ p_{i}(v^{S}) = \left( \prod_{k=j}^{j+m-1} q_{k}(v^{T}_{k}) \right) \cdot \pi_{j+m}(v^{S}_{suf}) \]
2.2.4. Special Case 3: Whitespace & Tabs Only
When a student’s token contains only whitespaces or tabs, BPM doesn’t provide a teacher’s reference probability. We simply ignore such tokens, to prevent code indentation errors.
2.2.5. Special Case 4: Stop Tokens and Bridging
The issue related to stop tokens is that the student’s EOS3 may not equal to the teacher’s EOS token. To make students correctly learn the signal to stop, we need a bridging target to describe the probability to stop.
In practice, we collect EOS tokens claimed in configuration, chat template, and generation config from the student and the teacher, and then remove them from the byte stream.
In addition, we extend the token mapping function \(\phi(\cdot)\) such that it maps the teacher’s stop token to the student’s stop token. So the stop tokens can join the training loss.
The bridging target is defined as:
\[ \mathrm{stop}^{T}_{i}(v^{S}_{stop}) = \sum_{v^{T}_{stop} \in \mathcal{V}^{T}_{stop}} q_{T(i)}(v^{T}_{stop}), \qquad \mathrm{stop}^{T}_{i}(\varnothing) = 1 - \mathrm{stop}^{T}_{i}(v^{S}_{stop}) \]
It’s a binary classification, corresponding to stop or not.
Student Projection. In order to distill to let student learn whether to stop, we also need to binary-classify on student’s stop-or-not distribution. It’s simply defined as
\[ \bar{\mathrm{stop}}^{S}_{i}(v^{S}_{stop}) = p_{i}(v^{S}_{stop}), \bar{\mathrm{stop}}^{S}_{i} = 1-\bar{\mathrm{stop}}^{S}_{stop} \]
2.3. Training Objective
The paper uses a generalized Jensen-Shannon divergence. Suppose the teacher’s (predicted) vocabulary distribution is \(D^{T}\) and the student’s \(D^{S}\), then for \(0 \lt \beta \lt 1\), it’s defined as
\[ D_{\beta}(D^{T} \parallel D^{S}) = \beta \mathrm{KL}(D^{T} \parallel m_{\beta}) + (1-\beta)\mathrm{KL}(D^{S} \parallel m_{beta}) \]
where \(m_{\beta}\) is the mixture of teacher’s and student’s distribution:
\[ m_{\beta} = \beta D^{T} + (1-\beta) D^{S} \]
When \(\beta=0\), \(D_{0} = \mathrm{KL}(D^{T} \parallel D^{S})\) is pure forward KL; when \(\beta=1\), we define \(D_{1}(D^{T}\parallel D^{S}) = \mathrm{KL}(D^{S} \parallel (1-\lambda)D^{T} + \lambda D^{S})\), here \(\lambda=0.1\).
3. Findings
Indeed, base case plus special case 1 have already accounted for \(\gt 99\%\) of the cases.